import cv2
import numpy as np
import onnxruntime as ort
import os
from typing import Tuple

# 配置参数（对应 C++ 中的 config 和常量）
INPUT_H = 48
INPUT_W = 168
OUTPUT_SIZE1 = 21 * 78  # 字符序列输出 (21 个位置，每个 78 类)
OUTPUT_SIZE2 = 5        # 颜色分类输出 (5 种颜色)

# 车牌字符集（注意：索引0通常表示空白或忽略）
plate_chr = [
    "<PAD>", "京", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑",
    "苏", "浙", "皖", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤", "桂",
    "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "港",
    "澳", "使", "学", "警", "挂", "军", "空", "北", "南", "广", "沈",
    "2",      # index 43: 表示数字 '2'
    "成", "济",
    "5",      # index 47: 表示数字 '5'
    "民", "航",
    "C",      # index 51: 表示字母 'A'
    "D",      # index 52: 表示字母 'B'
    "A", "B", "E", "F", "G", "H",
    "J", "K", "L", "M", "N", "P", "Q",
    "T",      # index 66: 表示字母 'R'
    "S", "R", "U", "V", "W", "X", "Y", "Z"
]

# 车牌颜色类别（与模型输出维度对应）
plate_color_list = ["blue", "green", "yellow", "white", "black"]

# 归一化参数（来自 C++ 代码中减去均值除以标准差）
MEAN = 0.588
STD = 0.193


class PlateRecognition:
    def __init__(self, model_path: str, gpu_id: int = 0, batch_size: int = 1):
        self.model_path = model_path
        self.batch_size = batch_size
        self.gpu_id = gpu_id

        # 检查 ONNX 模型是否存在
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(f"ONNX model not found at {model_path}")

        # 设置 ONNX Runtime 的执行提供者（优先使用 CUDA）
        providers = []
        if gpu_id >= 0 and 'CUDAExecutionProvider' in ort.get_available_providers():
            providers.append(('CUDAExecutionProvider', {'device_id': gpu_id}))
        else:
            print("Warning: CUDA not available, falling back to CPU.")
            providers.append('CPUExecutionProvider')

        # 加载 ONNX 模型
        self.session = ort.InferenceSession(self.model_path, providers=providers)

        # 获取输入输出名称（对应 C++ 中的 INPUT_BLOB_NAME, OUTPUT_BLOB_NAME1/2）
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [out.name for out in self.session.get_outputs()]
        assert len(self.output_names) == 2, "Model should have 2 outputs"

        # 假设输出顺序：output_1（字符）, output_2（颜色）
        # 根据名字判断哪个是哪个
        self.output_name_chars = None
        self.output_name_color = None
        for name in self.output_names:
            if "output_1" in name or "char" in name.lower():
                self.output_name_chars = name
            elif "output_2" in name or "color" in name.lower():
                self.output_name_color = name
        assert self.output_name_chars is not None and self.output_name_color is not None, \
            "Could not identify output names for characters and color"

        print(f"Loaded ONNX model: {self.model_path}")
        print(f"Input: {self.input_name}, Outputs: {self.output_names}")

    def preprocess(self, img: np.ndarray) -> np.ndarray:
        """
        预处理图像：resize -> HWC to CHW -> 归一化 ((pixel / 255) - mean) / std
        """
        if img.size == 0:
            raise ValueError("Empty input image")

        # 调整大小
        resized = cv2.resize(img, (INPUT_W, INPUT_H))  # (168, 48)

        # 转换为 float32 并归一化
        blob = resized.astype(np.float32) / 255.0
        blob = (blob - MEAN) / STD  # 使用全局均值和标准差

        # HWC -> CHW
        blob = blob.transpose(2, 0, 1)  # (3, 48, 168)

        # 扩展 batch 维度
        blob = np.expand_dims(blob, axis=0).astype(np.float32)  # (1, 3, 48, 168)

        return blob

    def postprocess_color(self, color_output: np.ndarray) -> str:
        """
        后处理颜色输出：取 argmax 得到最可能的颜色
        """
        # color_output shape: (batch_size, 5)
        color_probs = color_output[0]  # 取第一个样本
        max_idx = np.argmax(color_probs)
        return plate_color_list[max_idx]

    def postprocess_plate(self, char_output: np.ndarray) -> str:
        """
        后处理字符输出：每行取 argmax，然后去重去空白，拼接成字符串
        char_output shape: (batch_size, 21*78) -> reshape to (21, 78)
        """
        char_probs = char_output[0]  # (21*78,)
        char_probs = char_probs.reshape(21, 78)  # (21, 78)

        plate_index = []
        for i in range(21):
            max_idx = np.argmax(char_probs[i])
            plate_index.append(max_idx)

        # CTC-like 解码：去除重复 + 去除 0（<PAD>）
        plate_str = ""
        prev = -1
        for idx in plate_index:
            if idx != 0 and idx != prev:  # 忽略空白和重复
                if 0 <= idx < len(plate_chr):
                    plate_str += plate_chr[idx]
            prev = idx

        return plate_str

    def run(self, img: np.ndarray) -> Tuple[str, str]:
        """
        主推理函数
        :param img: BGR 图像 (cv2 格式)
        :return: (plate_str, plate_color)
        """
        # 预处理
        blob = self.preprocess(img)

        # 推理
        result = self.session.run(self.output_names, {self.input_name: blob})
        output_chars, output_color = result[0], result[1]

        # 后处理
        plate_str = self.postprocess_plate(output_chars)
        plate_color = self.postprocess_color(output_color)

        return plate_str, plate_color

    def release(self):
        """释放资源（Python 版本主要是清理 session）"""
        del self.session


# =========================
# 使用示例
# =========================
if __name__ == "__main__":
    # 示例配置
    config = {
        "PlateReconitionModelPath": r"D:\CodeCNN\PlateRecognition\PlateRecognition\test\plate_recognition_color.onnx",
        "gpu_id": 0,
        "plate_recognition_bs": 1
    }

    # 初始化识别器
    recognizer = PlateRecognition(
        model_path=config["PlateReconitionModelPath"],
        gpu_id=config["gpu_id"],
        batch_size=config["plate_recognition_bs"]
    )

    # 读取测试图像
    img_path = "cropped_20250906111554555_092240102/cropped_traffic_light_0.jpg"  # 替换为你的测试图片路径
    image = cv2.imread(img_path)
    if image is None:
        raise FileNotFoundError(f"Cannot load image from {img_path}")

    # 执行识别
    plate_str, plate_color = recognizer.run(image)

    print(f"Detected Plate: {plate_str}")
    print(f"Plate Color: {plate_color}")

    # 释放资源
    recognizer.release()